#!/usr/bin/python
import os, sys, time, zestlib as zest
from zestlib import which
################################################################################
################################################################################
def PrintHelpMsgAndDie():
    ( p, cmd ) = os.path.split( sys.argv[0] )
    print 'usage:', cmd, '<num_mem_images> <interval_in_seconds>'
    print '... num_mem_images      : integer specifiying number of memory images to capture, default = 3'
    print '... interval_in_seconds : real number specifying interval between image captures, default = 60'
    print ''
    sys.exit( -1 )

def ParseArgv():
    if len( sys.argv ) == 1:
        numimages = 1
        interval  = 0.0
        return ( numimages, interval )
    if '-h' in sys.argv or '--help' in sys.argv:
        PrintHelpMsgAndDie()
    try:
        numimages = int( sys.argv[1] )
        interval  = 0.0
        if numimages > 1:
            interval = float( sys.argv[2] )
            assert interval  > 0.0
        assert numimages > 0
    except:
        PrintHelpMsgAndDie()
    return ( numimages, interval )

################################################################################
################################################################################
def SanityChecks( imagepath, numimages, kmod, verbose = False ):
    # ensure the path exists
    zest.makedirp( imagepath )
    #
    assert os.path.isdir( imagepath )
    assert os.path.isfile( kmod )
    #
    totmem = zest.TotalMemAvail()
    diskneed = numimages * totmem
    diskfree = zest.df( imagepath )
    if verbose or diskneed > diskfree:
        print 'diskfree = %.2f GB' % ( float( diskfree ) / 1024.0 / 1024.0 / 1024.0 )
        print 'diskneed = %.2f GB' % ( float( diskneed ) / 1024.0 / 1024.0 / 1024.0 )
    if diskneed > diskfree:
        print ''
        print 'not enough disk space'
        sys.exit( -1 )
    assert diskneed < diskfree

################################################################################
################################################################################
def Virtualized( zestroot ):
    virtbin = os.path.join( zestroot, 'zest', 'bin', 'virtualized' )
    lines = zest.SystemWrapperPipelined( virtbin )
    inside_vm = False
    for line in lines:
        print line
        toks = line.split( '=' )
        assert len( toks ) == 2
        k = toks[0]
        v = toks[1]
        if k == 'virtualized':
            if v == 'true':
                inside_vm = True
    return inside_vm

def RdOk( p ):
    return os.access( p, os.R_OK )

def WrOk( p ):
    return os.access( p, os.W_OK )

def CheckMSRRdWr():
    msr0path = '/dev/cpu/0/msr'
    rdok = RdOk( msr0path )
    wrok = WrOk( msr0path )
    rdwrok = rdok and wrok
    return rdwrok

def PrintSysInfo( zestroot, msr_rd_wr ):
    os.environ[ 'PATH' ] += os.pathsep + os.sep + os.path.join( 'usr', 'sbin' )
    #
    freepath = which( 'free' )
    unampath = which( 'uname' )
    dmidpath = which( 'dmidecode' )
    #
    meminfo = RdOk( '/proc/meminfo' )
    cpuinfo = RdOk( '/proc/cpuinfo' )
    runfree = freepath is not None
    rundmid = dmidpath is not None
    rununam = unampath is not None
    #
    if runfree: zest.TrySystemWrapper( 'free', verbose = True )
    if meminfo: zest.TrySystemWrapper( 'cat /proc/meminfo', verbose = True )
    if cpuinfo: zest.TrySystemWrapper( 'cat /proc/cpuinfo', verbose = True )
    if rundmid: zest.TrySystemWrapper( 'dmidecode --type=17', verbose = True )
    if rununam: zest.TrySystemWrapper( 'uname -a', verbose = True )
    # 
    likwidtopo = os.path.join( zestroot, 'likwid-3.0', 'likwid-topology' )
    likwidfeat = os.path.join( zestroot, 'likwid-3.0', 'likwid-features' )
    zest.TrySystemWrapper( likwidtopo, verbose = True )
    if msr_rd_wr:
        zest.TrySystemWrapper( likwidfeat, verbose = True )

################################################################################
################################################################################
def IAPerf( zestroot ):
    interval = 10 # 10 seconds between stats capture
    number   = 3  # do 3 captures
    #
    iaperf = os.path.join( zestroot, 'iaperf', 'bin', 'pcm.x' )
    assert os.path.isfile( iaperf )
    cmd = ' '.join([ iaperf, str( interval ), str( number ) ])
    zest.TrySystemWrapper( cmd )

def FindValidLikwidGroups( likwid ):
    assert os.path.isfile( likwid )
    cmd = ' '.join([ likwid, '-a' ])
    lines = zest.TrySystemWrapperPipelined( cmd )
    valid = []
    for line in lines:
        line = line.rstrip()
        line = line.lstrip()
        toks = line.split( ':' )
        if len( toks ) == 2:
            valid.append( toks[0] )
    #
    # find the intersection of what is available and what we want:
    using   = []
    desired = [ 'L3', 'MEM', 'ENERGY' ]
    for g in desired:
        if g in valid:
            using.append( g )
    return using

def Likwid( zestroot, ncpus ):
    likwid = os.path.join( zestroot, 'likwid-3.0', 'likwid-perfctr' )
    groups = FindValidLikwidGroups( likwid )
    c = '-'.join([ '0', str( ncpus-1 ) ])
    s = str( 10 )
    for lidx in range( 3 ):
        for g in groups:
            group = ' '.join([ '-g', g ])
            cores = ' '.join([ '-c', c ])
            steth = ' '.join([ '-S', s ])
            lcmd  = ' '.join([ likwid, group, cores, steth ])
            zest.TrySystemWrapper( lcmd )

def Lime( imagepath, imgfiles ):
    fnts  = zest.FileNameTimeStamp()
    iname  = '.'.join([ 'memimage', fnts ])
    rcname = '.'.join([ 'pagerefs', fnts ])
    iname  = os.path.join( imagepath,  iname )
    rcname = os.path.join( imagepath, rcname )
    patharg   = '='.join([ 'path',   iname  ])
    rcpatharg = '='.join([ 'rcpath', rcname ])
    icmd = ' '.join([ '/sbin/insmod', kmod, patharg, rcpatharg, 'format=raw' ])
    rcmd = ' '.join([ '/sbin/rmmod', 'lime' ])
    zest.SystemWrapper( icmd )
    zest.SystemWrapper( rcmd )
    imgfiles.append( iname )

################################################################################
################################################################################
verbose = True
( numimages, interval ) = ParseArgv()

# verify user is root, grab # of cpus, & grab relevant paths:
zest.VerifyRoot()
ncpus = os.sysconf( 'SC_NPROCESSORS_ONLN' )
zestroot  = zest.ZestRoot()
imagepath = zest.NewImagePath()
inside_vm = Virtualized( zestroot )
msr_rd_wr = CheckMSRRdWr()
likwidiap = msr_rd_wr and not inside_vm
kmod   = os.path.join( zestroot, 'lime', 'src', 'lime.ko' )

# setup the log file:
logfile = '.'.join([ 'zest', 'capture', zest.FileNameTimeStamp(), 'log' ])
logfile = os.path.join( imagepath, logfile )
zest.opentee( logfile )

SanityChecks( imagepath, numimages, kmod, verbose = False )
PrintSysInfo( zestroot, msr_rd_wr )

imgfiles = []

# the main loop:
# for each image:
# 1. profile with iaperf
# 2. profile with likwid
# 3. grab the mem image
for iidx in range( numimages ):
    if iidx:
        print '... waiting for', interval, 'seconds'
        time.sleep( interval )
    SanityChecks( imagepath, numimages - iidx, kmod, verbose )
    if likwidiap:
        # run likwid & ia-perf
        IAPerf( zestroot )
        Likwid( zestroot, ncpus )
    Lime( imagepath, imgfiles )
    

print '--------------------------------------------------------------------------------'
print '-- created image files:'
print '\n'.join( imgfiles )
print '-- output in log file:'
print logfile


